from PIL.Image import register_extensions
from openai import OpenAI
import json

# 初始化 OpenAI 客户端
client = OpenAI(
    base_url="http://localhost:11434/v1",  # Ollama API 地址
    api_key="sk-xxx"  # 可以是任意值
)


# 定义可用函数（示例：天气查询）
def get_weather(location):
    """获取指定城市的当前天气"""
    # 实际实现中可以调用真实天气 API
    weather_data = {
        "location": location,
        "temperature": "22°C",
        "condition": "晴朗",
        "humidity": "45%"
    }
    return json.dumps(weather_data)


# 定义函数列表（与 MCP 服务端的 tools 对应）
# 修改工具定义格式
tools = [
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "获取指定城市的当前天气",
            "parameters": {
                "type": "object",
                "properties": {
                    "location": {
                        "type": "string",
                        "description": "城市名称，如北京、上海"
                    }
                },
                "required": ["location"]
            }
        }
    }
]


# 调用模型并处理函数调用
def call_model_with_functions(prompt):
    response = client.chat.completions.create(
        model="MFDoom/deepseek-r1-tool-calling:14b",
        messages=[{"role": "user", "content": prompt}],
        tools=tools,
        tool_choice={"type": "function", "function": {"name": "get_weather"}}
    )

    message = response.choices[0].message
    print("模型原始响应:", message)  # 添加调试输出

    if message.tool_calls:
        for tool_call in message.tool_calls:
            if tool_call.function.name == "get_weather":
                args = json.loads(tool_call.function.arguments)
                weather = get_weather(args["location"])

                # 发送函数执行结果回模型
                response = client.chat.completions.create(
                    model="deepseek:7b-chat",
                    messages=[
                        {"role": "user", "content": prompt},
                        message,
                        {
                            "role": "tool",
                            "content": weather,
                            "tool_call_id": tool_call.id
                        }
                    ]
                )
                return response.choices[0].message.content
    else:
        print("警告：模型未按预期调用工具")
        # 可以在这里添加回退逻辑或重新提示

    return message.content

user_query = "上海今天天气如何？"
response = call_model_with_functions(user_query)
print(f"回答：{response}")


#定义需要显示 的文字及其格式
text = "hello world"
response = call_model_with_functions(text)

# 输出示例：
# 回答：上海当前的天气为晴朗，温度22°C，湿度45%。